#!/bin/sh


function GetVPCID() {
	string=$(curl -s --max-time 1 100.100.100.200/latest/meta-data/vpc-id)
	uppercase=$(echo "$string" | tr '[:lower:]' '[:upper:]')
	echo "$uppercase"
}

# check if command exist, 0 for yes, other for no
function CheckCommand() {
	type "$1" > /dev/null 2>&1
}

function err() {
	echo "$@"
} > /dev/stderr

# 0 for success, other for failed
function CheckDependencies() {

	CheckCommand eadm
	if [ $? != 0 ]; then
		err "no eadm found."
		return 1;
	fi

	CheckCommand smcr
	if [ $? != 0 ]; then
		err "no smcr found."
		return 2;
	fi

	CheckCommand bpftool
	if [ $? != 0 ]; then
		err "no bpftool found."
		return 3;
	fi

	CheckCommand smc-ebpf
	if [ $? != 0 ]; then
		err "no smc-ebpf found."
		return 4;
	fi

	return 0;
}

function CheckUEID() {
	arr=($(smcr ueid))
	vpcid=$(GetVPCID)
	if [ ${#arr[@]} -eq 1 ] && [ x"${arr[0]}" == x"$vpcid" ]; then
		return 0;
	fi
	err "smc UEID not set correctly"
	return 5;
}

# 0 for success, other for failed
function CheckPermissions() {

	if [ $(id -u) == 0 ]; then
		return 0;
	fi

	err "no permission"
	return 6;
}

# 0 for success, other for failed
function checkDevice() {
	ibv_devinfo > /dev/null 2>&1
	if [ $? != 0 ]; then
		err "no eRDMA device"
		return 7;
	fi
	return 0;
}

# 0 for success, other for failed
function checkSMC() {

	modinfo smc > /dev/null 2>&1
	# no smc module
	if [ $? != 0 ]; then
		err no smc modules
		return 8;
	fi

	# required kernel version greater than 5.10.134-16
	KERNEL_VER=$(uname -r)
	REQUIRED_VER="5.10.134-16"
	if [[ "$(printf '%s\n' "$KERNEL_VER" "$REQUIRED_VER" | sort -V | head -n1)" != "${REQUIRED_VER}" ]]; then
		err "kernel version not match"
		return 9;
	fi

	return 0;
}

function ConverNumber() {

	str="$1"

	if [[ $str == 0b* ]]; then
		echo $((2#${str#0b}))
	elif [[ $str == 0o* ]]; then
		echo $((8#${str#0o}))
	elif [[ $str == 0x* ]]; then
		echo $((16#${str#0x}))
	else
		echo $str
	fi
}

function CheckFeature() {
	# check if IM support
	output=`eadm info -t ext_attr | grep ^ext_cap_mask`
	mask=`echo $output | grep -oE '(0x[0-9a-fA-F]+|0[0-7]+|0b[01]+|[1-9][0-9]*)'`
	value=$(ConverNumber "$mask")
	if [ x"$value" == x ]; then
		err requires IM support
		return 10;
	fi
	if [[ $(($value & 3)) -ne 3 ]]; then
		err "requires IM support ($value)"
		return 11;
	fi
	return 0;
}

function hasERDMA() {

	# first try install dependencies
	initERDMA

	# check package dependencies
	CheckDependencies
	rc=$?

	if [ ${rc} != 0 ]; then
		return ${rc};
	fi

	# check permission to on/off smc
	CheckPermissions
	rc=$?

	if [ ${rc} != 0 ]; then
		return ${rc};
	fi

	# check if there are eRDMA device
	checkDevice
	rc=$?

	if [ ${rc} != 0 ]; then
		return ${rc};
	fi

	# check if there are SMCv2.1 support
	checkSMC
	rc=$?

	if [ ${rc} != 0 ]; then
		return ${rc};
	fi

	# check if SMCv2.1 and IM support)
	CheckFeature
	rc=$?

	if [ ${rc} != 0 ]; then
		return ${rc};
	fi

	return 0
}

function isERDMA() {
	ret=`sysctl -q net.smc.tcp2smc | awk -F= '{print $2}' | xargs`
	if [ x"$ret" == x1 ]; then
		echo on
		return
	fi
	echo off
}

function removeUEID() {
	smcr ueid | grep $1
	if [ $? == 0 ]; then
		smcr ueid del $1
	fi
} > /dev/null 2>&1

function addUEID() {
	smcr ueid | grep $1
	if [ $? != 0 ]; then
		smcr ueid add $1
	fi
} > /dev/null 2>&1

function getERDMA_DEVNAME() {
	ibv_devinfo | grep hca_id | awk  '{print $2}'
}

function enableERDMA() {
	smc-ebpf policy load --init
	if [ $? != 0 ]; then
		return 13;
	fi

	eadm conf -t  dack_count -d $(getERDMA_DEVNAME) -v 5
	if [ $? != 0 ]; then
		return 14;
	fi

	sysctl -w net.smc.tcp2smc=1;
} > /dev/null 2>&1

function disableERDMA() {
	sysctl -w net.smc.tcp2smc=0;
	eadm conf -t  dack_count -d $(getERDMA_DEVNAME) -v 3;
	smc-ebpf policy clear
} > /dev/null 2>&1

function toggleERDMA() {

	current=$(isERDMA)
	target=$(echo $1 | tr '[:upper:]' '[:lower:]')

	# already set
	if [ x"$target" == x"$current" ]; then
		return 0;
	fi

	if [ x"$target" == x'on' ]; then
		modprobe smc
		if [ $? != 0 ]; then
			err "insert smc module fail."
			return 12;
		fi

		# change default memory configure
		sysctl -w net.smc.wmem=65536
		sysctl -w net.smc.rmem=65536

		removeUEID SMCV2-DEFAULT-UEID
		addUEID $(GetVPCID)

		enableERDMA
		ret=$?

		if [ x"$ret" != x0 ]; then
			err "enableERDMA fail by ${ret}"
			return $ret
		fi
	else # off
		disableERDMA
	fi

	return 0;
}

function initERDMA()
{
	CheckCommand smcr
	if [ $? != 0 ]; then
		yum install -y smc-tools
	fi

	CheckCommand bpftool
	if [ $? != 0 ]; then
		yum install -y bpftool
	fi

	# add repo for eRDMA to install eadm
	yum-config-manager  --add-repo http://mirrors.cloud.aliyuncs.com/erdma/yum/alinux/erdma.repo

	yum install -y eadm libibverbs-utils

} > /dev/null 2>&1

function="$1"
shift
${function} $@
